package com.ask.forMe.controller.ai;

import com.ask.forMe.langchain4j.agent.DocxMarkHeadingAgent;
import com.ask.forMe.langchain4j.model.entity.Section;
import com.ask.forMe.langchain4j.model.vo.SplitterVO;
import com.ask.forMe.langchain4j.rag.loader.MyDocumentLoader;
import com.ask.forMe.langchain4j.rag.splitter.LangchainSplitter;
import com.ask.forMe.langchain4j.rag.splitter.RegexSplitUtils;
import com.ask.forMe.langchain4j.rag.splitter.SplitMdByHeadingUtil;
import com.ask.forMe.model.entity.Result;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.segment.TextSegment;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PutMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

import static com.ask.forMe.constant.FileConstant.BASE_PATH;
import static com.ask.forMe.constant.SystemConstant.KNOWLEDGE_PATH;

@RestController
@RequestMapping("/split")
@Tag(name = "Splitter接口")
@Slf4j
public class SplitterController {

    @Autowired
    private DocxMarkHeadingAgent docxMarkHeadingAgent;

    @PutMapping("/split")
    @Operation(summary = "split方法测试")
    public Result testSplitFunctions(@RequestBody SplitterVO splitVo) {
        Document document = MyDocumentLoader.loadAndParserByApacheTika(BASE_PATH + splitVo.getFileName());
        String[] split = RegexSplitUtils
                .splitByRegex(document.text(), splitVo.getRegex());
        System.out.println("分段数：" + split.length);
        for (String s : split) {
            if (!s.trim().isEmpty()) {
                System.out.println("【----------------分段----------------】");
                System.out.println(s);
            }
        }
        return Result.success(split);
    }

    @PutMapping("/splitter")
    @Operation(summary = "splitter测试")
    public Result testSplitter(@RequestBody SplitterVO splitVo,String type) {
        // 读取文档，使用分割器分块
        Document document = MyDocumentLoader.loadAndParserByApacheTika(BASE_PATH + splitVo.getFileName());
        List<TextSegment> segments = LangchainSplitter.splitterSelector(type, document.text(), splitVo);
        // 输出分割后的内容
        List<String> res = segments.stream()
                .filter(segment -> segment != null && !segment.text().trim().isEmpty())
                .map(TextSegment::text)
                .toList();
        res.forEach(r -> {
            System.out.println("\n【----------------分段----------------】");
            System.out.println(r);
        });
        return Result.success(res);
    }

    @PutMapping("/splitDocxByFullHeading")
    @Operation(summary = "分割文档-完整标题路径")
    public Result splitWordByFullHeading(@RequestBody SplitterVO splitVo) {
        List<Section> sections = SplitMdByHeadingUtil.
                splitMDWithFullHeadingPath(splitVo.getFileName(), splitVo.getMaxSegmentSizeInChars(), splitVo.getMaxHeadingLevel());
        sections.forEach(section -> {
            System.out.println("\n【----------------分段----------------】");
            System.out.println("level: " + section.getLevel());
            System.out.println("titlePath: " + section.getTitlePath());
            System.out.println("content: " + section.getContent());
        });
        return Result.success(sections);
    }

    @PutMapping("/splitDocxByHeading")
    @Operation(summary = "按标题分割-LLM标记")
    public Result splitWordByHeading(@RequestBody SplitterVO splitVo) {
        Document document = MyDocumentLoader.loadAndParserByApacheTika(splitVo.getFilePath());
        // LLM为标题增加标记
        String markedText = docxMarkHeadingAgent.chat(document.text());
        // 根据标题拆分
        String[] splitByHeading = RegexSplitUtils.splitByHeading(markedText, splitVo.getMinHeadingLevel(), splitVo.getMaxHeadingLevel());
        return Result.success(splitByHeading);
    }
}
